from manim import *
from pi_creature_scene import (
    PiCreatureScene,
    MortyPiCreatureScene,
    TeacherStudentsScene,
)


class BasicPiCreatureScene(PiCreatureScene):
    def construct(self):
        # 基本表情变换演示
        self.pi_creature.shift(DOWN)
        # 确保初始时 pi 生物在场景中
        if self.pi_creature not in self.mobjects:
            self.add(self.pi_creature)

        modes = ["plain", "happy", "confused", "thinking", "surprised", "sassy"]
        mode_texts = [Text(mode, font_size=24).to_corner(UP + LEFT) for mode in modes]

        self.play(Write(mode_texts[0]))
        for i, mode in enumerate(modes[1:]):
            self.play(
                self.pi_creature.animate.change_mode(mode),
                ReplacementTransform(mode_texts[i], mode_texts[i + 1]),
            )
            self.wait()

        # 对话和思考
        self.play(FadeOut(mode_texts[-1]))
        self.say("我可以说话!")
        self.wait()
        self.think("我也可以思考...")
        self.wait()

        # 目光追踪
        dot = Dot(color=YELLOW).to_edge(RIGHT, buff=2)
        self.play(Create(dot))
        self.play(
            dot.animate.to_edge(LEFT, buff=2), self.pi_creature.animate.look_at(dot)
        )
        self.play(
            dot.animate.to_edge(UP, buff=2), self.pi_creature.animate.look_at(dot)
        )
        self.play(FadeOut(dot))

        # 眨眼演示前再次确认
        if self.pi_creature not in self.mobjects:
            self.add(self.pi_creature)

        for _ in range(3):
            self.blink()
            self.wait()

        self.pi_creature_says(
            self.pi_creature, "这就是\nPiCreatureScene!", target_mode="happy"
        )
        self.wait(2)


class MortyScene(MortyPiCreatureScene):
    def construct(self):
        # Morty的特点演示
        self.play(self.pi_creature.animate.shift(RIGHT))
        formula = MathTex("e^{i\\pi} + 1 = 0")

        # Morty反应
        self.play(Write(formula))
        self.play(self.pi_creature.animate.change_mode("thinking"))
        self.wait()
        self.play(self.pi_creature.animate.change_mode("surprised"))
        self.wait()

        # 对话
        self.pi_creature_says(
            self.pi_creature, "这个公式\n太美了!", target_mode="happy"
        )
        self.wait(2)

        # 眼神变化
        box = SurroundingRectangle(formula)
        self.play(Create(box))
        self.play(self.pi_creature.animate.look_at(formula))
        self.wait()

        # 展示思考过程
        self.pi_creature_thinks(
            self.pi_creature,
            "它连接了数学中\n五个最重要的常数...",
            target_mode="thinking",
        )
        self.wait(2)


class TeacherStudentDemo(TeacherStudentsScene):
    def construct(self):
        # 基本教学设置
        self.teacher_says("让我们来学习\nPi生物的特性", target_mode="speaking")
        self.play_all_student_changes("confused")
        self.wait()

        # 学生提问
        self.student_says(
            "Pi生物有\n多少种表情?", target_mode="raise_right_hand", index=1
        )
        self.play(self.teacher.animate.change_mode("thinking"))
        self.wait()

        # 老师回答
        self.teacher_says(
            "有很多种!\n比如: happy, plain, confused...", target_mode="happy"
        )
        self.play_student_changes("happy", "thinking", "happy")
        self.wait()

        # 老师展示内容
        formula = MathTex("\\pi r^2").scale(1.5)
        self.teacher_holds_up(formula)
        self.play_student_changes("surprised", "surprised", "surprised")
        self.wait()

        # 学生思考 - 直接使用动画结果
        thought_anim = self.student_thinks(
            "这真是\n太有趣了!", target_mode="thinking", index=2
        )

        # 重要：获取学生和思考泡泡
        student = self.get_students()[2]
        print(f"学生: {student}, 有泡泡: {hasattr(student, 'bubble')}")

        # 不要再调用wait()，这可能会移除泡泡
        # 直接放大
        if hasattr(student, "bubble") and student.bubble is not None:
            self.zoom_in_on_thought_bubble(bubble=student.bubble)
            self.wait(2)
        else:
            # 调试信息 - 如果找不到泡泡
            print("警告：找不到思考泡泡！")
            self.wait(2)


class AllSceneDemo(Scene):
    def construct(self):
        # 标题
        title = Text("Pi生物场景演示", font_size=56)
        self.play(Write(title))
        self.wait()
        self.play(FadeOut(title))

        # 提示：每个场景需要单独渲染
        text = Text("请分别渲染各个场景：", font_size=36).shift(UP)
        cmds = (
            VGroup(
                Text(
                    "manim -pqm pi_creature_demo.py BasicPiCreatureScene", font_size=24
                ),
                Text("manim -pqm pi_creature_demo.py MortyScene", font_size=24),
                Text("manim -pqm pi_creature_demo.py TeacherStudentDemo", font_size=24),
            )
            .arrange(DOWN, aligned_edge=LEFT)
            .next_to(text, DOWN, buff=0.5)
        )

        self.play(Write(text))
        self.play(Write(cmds))
        self.wait(3)


if __name__ == "__main__":
    from os import system

    # 直接渲染单个场景
    system("manim -pqm  {}   --disable_caching  TeacherStudentDemo".format(__file__))
